import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS

def model(data):
    reversion_speed = numpyro.sample("k",dist.HalfNormal(1.0))
    sigma = numpyro.sample("sigma",dist.HalfNormal(1.0))
    x0 = numpyro.sample(
        "x0",
        dist.Normal(0.0, jnp.divide(jnp.power(sigma, 2), 2 * reversion_speed)),
        obs=data[0]
    )

    numpyro.sample(
        "x",
        dist.GaussianStateSpace(
            num_steps=jnp.shape(data)[-1]-1,
            transition_matrix=jnp.array([[ jnp.exp(- jnp.multiply(reversion_speed, 1.0)) - 1 ]]),
            covariance_matrix=jnp.array([[ jnp.multiply(jnp.divide(jnp.power(sigma, 2), 2), jnp.divide(1 - jnp.exp(- jnp.multiply(reversion_speed, 1.0)), reversion_speed)) ]]),
            initial_value = jnp.stack([x0])
        ),
        obs=jnp.stack([data[1:]], axis=-1)
    )

mcmc = MCMC(
    NUTS(
        model=model
    ), 
    num_warmup=500, 
    num_samples=1_000
)
mcmc.run(
    rng_key=jax.random.PRNGKey(2),
    data=(
        lambda reversion_speed, sigma, std_norm: jnp.concatenate([
            jnp.array([jnp.sqrt(jnp.divide(jnp.power(sigma, 2), 2 * reversion_speed)) * std_norm[0]]),
            jax.lax.scan(
                lambda y, x: (jnp.multiply(y, jnp.exp(- jnp.multiply(reversion_speed, 1.0)) - 1) + x * jnp.sqrt(jnp.multiply(jnp.divide(jnp.power(sigma, 2), 2), jnp.divide(1 - jnp.exp(- jnp.multiply(reversion_speed, 1.0)), reversion_speed))),) * 2,
                init=jnp.sqrt(jnp.divide(jnp.power(sigma, 2), 2 * reversion_speed)) * std_norm[0],
                xs=std_norm[1:]
            )[1]
        ])
    )(
        0.1, 0.5, jax.random.normal(key=jax.random.PRNGKey(10), shape=(20,))
    )
)
mcmc.print_summary()